#include "pytorch_cpp_helper.hpp"
Tensor NMSCUDAKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold, int offset);
Tensor nms_cpu(Tensor boxes,Tensor scores,float iou_threshold,int offset)
{
    /*
        基本思路：
        (1) 取出x1,y1,x2,y2,areas_t,scores
        (2) 循环将bbox和其他剩余的bbox进行交并比计算，将交并比大于阈值的bbox从这个集合中剔除出去,设为false
        (3) 继续循环
        为了提高效率，我们保留bbox不动，最终保留的也都是bbox在原集合中的索引(mask_select)
    */
    if(boxes.numel()==0){
        return at::empty({0},boxes.options().dtype(at::kLong));
    }
    // 获取boxes的各个维度
    auto x1_t =boxes.select(1,0).contiguous();
    auto y1_t = boxes.select(1,1).contiguous();
    auto x2_t = boxes.select(1,2).contiguous();
    auto y2_t = boxes.select(1,3).contiguous();
    // 获取每个box的面积
    Tensor areas_t = (x2_t - x1_t + offset) * (y2_t - y1_t + offset);
    auto order_t =std::get<1> (scores.sort(0,true));
    
    auto nboxes = boxes.size(0);
    Tensor select_t = at::ones({nboxes}, boxes.options().dtype(at::kBool));
    // 用data_ptr可以很方便的获取一个tensor的元素指针，从而访问tensor
    auto select = select_t.data_ptr<bool>();
    auto order = order_t.data_ptr<int64_t>();
    auto x1 = x1_t.data_ptr<float>();
    auto y1 = y1_t.data_ptr<float>();
    auto x2 = x2_t.data_ptr<float>();
    auto y2 = y2_t.data_ptr<float>();
    auto areas = areas_t.data_ptr<float>();

    for (int64_t _i = 0; _i < nboxes; _i++) {
        if (select[_i] == false) continue;
        auto i = order[_i];
        auto ix1 = x1[i];
        auto iy1 = y1[i];
        auto ix2 = x2[i];
        auto iy2 = y2[i];
        auto iarea = areas[i];

        for (int64_t _j = _i + 1; _j < nboxes; _j++) {
            if (select[_j] == false) continue;
            auto j = order[_j];
            auto xx1 = std::max(ix1, x1[j]);
            auto yy1 = std::max(iy1, y1[j]);
            auto xx2 = std::min(ix2, x2[j]);
            auto yy2 = std::min(iy2, y2[j]);

            auto w = std::max(0.f, xx2 - xx1 + offset);
            auto h = std::max(0.f, yy2 - yy1 + offset);
            auto inter = w * h;
            auto ovr = inter / (iarea + areas[j] - inter);
            if (ovr > iou_threshold) select[_j] = false;
        } 
    }
    return order_t.masked_select(select_t);
}


PYBIND11_MODULE(my_ops, m)
{
    m.def("nms", nms_cpu, "nms_compute",
        py::arg("boxes"), py::arg("scores"), py::arg("iou_threshold"),
        py::arg("offset"));
    m.def("nms_cuda", NMSCUDAKernelLauncher, "nms_compute_cuda",
        py::arg("boxes"), py::arg("scores"), py::arg("iou_threshold"),
        py::arg("offset"));
}